# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/5 11:11 
@Author  : ZhangShenao 
@File    : summarize_buffer_memory.py 
@Desc    : 自定义的摘要-缓冲记忆组件
"""
import os

from zhipuai import ZhipuAI

# 生成摘要的Prompt
SUMMARIZE_PROMPT = '''你是一个强大的聊天机器人，具有很强的摘要总结能力。
请根据用户提供的历史摘要信息和历史聊天记录，总结和提炼关键信息，生成摘要。要求摘要尽可能还原原始内容。
下面是一个示例，帮助你更好地理解任务的要求。

<示例开始>

【历史信息摘要】
人类会问人工智能对人工智能的看法。人工智能认为人工智能是一股向善的力量。

【历史聊天记录】
用户提问: 为什么你认为人工智能是一股向善的力量？
AI回答: 因为人工智能将帮助人类充分发挥潜力。

【生成的新摘要】
人类会问人工智能对人工智能的看法。人工智能认为人工智能是一股向善的力量，因为它将帮助人类充分发挥潜力。

</示例结束>

下面开始你的任务：

【历史信息摘要】
 {history_summarize}

 【历史聊天记录】
 {chat_history_msg}

【生成的新摘要】
'''

# 生成摘要的模型
SUMMARIZE_MODEL = 'glm-4-air'


class SummarizeBufferMemory:
    """自定义的摘要-缓冲记忆组件"""

    def __init__(self, max_tokens: int = 200):
        """
        构造函数
        :param max_tokens: 缓冲的最大token数量
        """

        self.__max_tokens = max_tokens  # 初始化缓冲的最大token数量
        self.__history_summarize = ''  # 初始化历史摘要
        self.__chat_history = []  # 初始化聊天历史列表
        self.__summarize_llm = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))  # 初始化摘要LLM客户端

    def save_chat_message(self, message: str):
        """
        将聊天消息保存到历史聊天列表
        :param message: 聊天消息，包含role和content字段
        """
        self.__chat_history.append(message)

        # 如果历史聊天列表的长度超过了max_token限制,则生成摘要,并将最历史最久的聊天记录从列表中移除
        history_msg = self.format_chat_history()
        token = self.calc_token_count(history_msg)
        while token > self.__max_tokens:
            # 根据历史摘要和历史聊天记录,生成最新的摘要
            print('新摘要生成中...')
            new_summarize = self.summarize(history_msg)
            print(f'生成的新摘要: {new_summarize}')

            # 更新历史摘要
            self.__history_summarize = new_summarize

            # 移除最历史最久的聊天记录
            self.__chat_history.pop(0)

            # 计算最新的历史聊天记录的token数量
            history_msg = self.format_chat_history()
            token = self.calc_token_count(history_msg)

    @classmethod
    def calc_token_count(cls, content: str) -> int:
        """
        计算内容包含的token数量
        :param content: 字符串内容
        :return: token数量
        """

        # TODO: 这里为了实现简便,暂时使用字符串长度作为token数量,实际应该调用大模型的token统计API
        return len(content)

    def format_chat_history(self) -> str:
        """
        格式化聊天历史列表为字符串
        :return: 格式化后的字符串
        """
        return '\n'.join(self.__chat_history)

    def summarize(self, chat_history_msg: str) -> str:
        """
        根据历史摘要和历史聊天记录,生成最新的摘要
        :return: 生成的新摘要
        """

        content = SUMMARIZE_PROMPT.format(history_summarize=self.__history_summarize, chat_history_msg=chat_history_msg)
        messages = [{'role': 'user', 'content': content}]
        response = self.__summarize_llm.chat.completions.create(messages=messages, model=SUMMARIZE_MODEL)
        return response.choices[0].message.content

    def get_history_summarize(self) -> str:
        """
        获取历史摘要
        :return: 历史摘要
       """
        return self.__history_summarize
